# -*- coding: utf-8 -*-
import functools
import os
import pickle
import time
import zipfile

import tensorflow as tf

import pandas as pd
from pyhive import hive


def connect_hive():
    # 建立连接
    hive_hotel_host = '192.168.7.24'
    hive_hotel_port = 10000
    hive_hotel_username = 'mfw_hadoop'
    hive_hotel_database = 'mdd'
    hive_hotel_auth = 'NOSASL'

    try:
        conn = hive.Connection(host=hive_hotel_host,
                               port=hive_hotel_port,
                               username=hive_hotel_username,
                               database=hive_hotel_database,
                               auth=hive_hotel_auth)
    except:
        conn = hive.Connection(host='192.168.7.23',
                               port=hive_hotel_port,
                               username=hive_hotel_username,
                               database=hive_hotel_database,
                               auth=hive_hotel_auth)

    return conn


def load_data_without_date(table_name):
    conn = connect_hive()
    data = pd.read_sql("""
    select * from %s
    """ % (table_name), conn)
    data.columns = [col.split(".")[-1] for col in data.columns]
    return data


def load_data_with_date(table_name, start, end):
    conn = connect_hive()
    data = pd.read_sql("""
    select * from %s
    where dt between '%s' and '%s'
    """ % (table_name, start, end), conn)
    data.columns = [col.split(".")[-1] for col in data.columns]
    return data


def load_rec_without_date(table_name):
    data = load_data_without_date("recommend." + table_name)
    return data


def load_rec_with_date(table_name, start, end):
    data = load_data_with_date("recommend." + table_name, start, end)
    return data


def unzip_file(filename, output_dir):
    z_file = zipfile.ZipFile(filename, 'r')
    os.makedirs(output_dir, exist_ok=True)
    for f in z_file.namelist():
        z_file.extract(f, output_dir)
        print(f)
    z_file.close()


def choose_gpu(is_choose=True, num=0):
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # Currently, memory growth needs to be the same across GPUs
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            if is_choose:
                tf.config.experimental.set_visible_devices(gpus[num], 'GPU')
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
        except RuntimeError as e:
            # Memory growth must be set before GPUs have been initialized
            print(e)


def save_by_pickle(obj, path, patten='wb'):
    with open(path, patten) as f:
        pickle.dump(obj, f, protocol=pickle.HIGHEST_PROTOCOL)


def load_pickle(path):
    with open(path, 'rb') as f:
        obj = pickle.load(f)
    return obj


def timeit(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        begin = time.time()
        result = fn(*args, **kwargs)
        print("cost %.2fs for {%s}" % (time.time() - begin, fn.__name__))
        return result

    return wrapper
